import streamlit as st
import akshare as ak
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from datetime import datetime

# 设置全局样式
plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows系统使用此字体
plt.rcParams['axes.unicode_minus'] = False

# 自定义双色系
dual_colors = ['#4CB050', '#FF3333']  # 绿跌红涨
cmap = mcolors.ListedColormap(dual_colors)
bounds = [-100, 0, 100]  # 严格分界
norm = mcolors.BoundaryNorm(bounds, cmap.N)

# 数据获取
@st.cache_data
def get_data():
    """获取上证指数历史数据"""
    df = ak.stock_zh_index_daily(symbol="sh000001")
    df['date'] = pd.to_datetime(df['date'])
    df['year'] = df['date'].dt.year
    df['month'] = df['date'].dt.month
    return df.sort_values('date')

# 构建月度数据表
def build_table(df, end_year):
    # 计算起始年份
    start_year = end_year - 10
    
    # 数据验证
    required_cols = ['year', 'month', 'close']
    if not all(col in df.columns for col in required_cols):
        missing = [col for col in required_cols if col not in df.columns]
        raise ValueError(f"缺少必要列: {missing}")

    # 筛选目标年份范围
    df = df[df['year'].between(start_year, end_year)]
    
    # 计算月度涨跌幅
    def calc_monthly_change(group):
        if len(group) >= 3:  # 至少3个交易日才计算
            return (group['close'].iloc[-1] / group['close'].iloc[0] - 1) * 100
        return None

    monthly_chg = (
        df.groupby(['year', 'month'])
        .apply(calc_monthly_change)
        .unstack(level='month')
        .rename(columns=lambda m: f"{int(m)}月")
        .round(2)
    )
    
    # 格式化输出
    result = (
        monthly_chg.loc[start_year:end_year]
        .reset_index()
        .rename(columns={'year': '年份'})
        .set_index('年份')
    )
    
    return result

# 生成热力图
def plot_heatmap(data):
    """生成月度涨跌幅热力图"""
    fig, ax = plt.subplots(figsize=(16, 6))
    sns.heatmap(data,
                cmap=cmap,
                norm=norm,
                annot=True,
                fmt=".1f",
                linewidths=0.5,
                cbar=False,
                annot_kws={'color': 'white', 'weight': 'bold'})
    ax.set_xticklabels(['1月', '2月', '3月', '4月', '5月', '6月',
                        '7月', '8月', '9月', '10月', '11月', '12月'])
    ax.set_title('月度涨跌幅分布', fontsize=14)
    return fig

def color_cell(val):
    """单元格颜色格式化"""
    if pd.isna(val):
        return ''
    color = '#FF3333' if val >= 0 else '#4CB050'
    return f'background-color: {color}; color: white'

# 主程序
def app():
    st.subheader("上证指数历年月度涨跌幅分析")
    
    # 获取当前年份
    current_year = datetime.now().year
    
    # 创建年份选择器
    selected_year = st.slider(
        "选择结束年份",
        min_value=2015,
        max_value=current_year,
        value=current_year,
        step=1
    )
    
    # 获取数据
    df = get_data()
    
    # 构建表格
    monthly_df = build_table(df, selected_year)
    
    # 显示表格
    st.write(f"上证指数{selected_year-10}-{selected_year}年每月涨跌幅数据")
    
    # 格式化并显示表格
    styled_df = monthly_df.style.format('{:.1f}%', na_rep="-").applymap(color_cell)
    st.dataframe(styled_df, height=450, use_container_width=True)
    
    # 显示热力图
    #st.pyplot(plot_heatmap(monthly_df))

if __name__ == "__main__":
    app()
